{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Matrix Multiplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-06-30 01:37:11.221 | DEBUG    | ttnn:<module>:83 - Initial ttnn.CONFIG:\n",
      "Config{cache_path=/home/ubuntu/.cache/ttnn,model_cache_path=/home/ubuntu/.cache/ttnn/models,tmp_dir=/tmp/ttnn,enable_model_cache=false,enable_fast_runtime_mode=true,throw_exception_on_fallback=false,enable_logging=false,enable_graph_report=false,enable_detailed_buffer_report=false,enable_detailed_tensor_report=false,enable_comparison_mode=false,comparison_mode_should_raise_exception=false,comparison_mode_pcc=0.9999,root_report_path=generated/ttnn/reports,report_name=std::nullopt,std::nullopt}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-06-30 01:37:11.927 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.928 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.946 | info     |          Device | Opening user mode device driver (tt_cluster.cpp:174)\n",
      "2025-06-30 01:37:11.947 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.947 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.951 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.952 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:11.956 | info     |   SiliconDriver | Harvesting mask for chip 0 is 0x40 (NOC0: 0x40, simulated harvesting mask: 0x0). (cluster.cpp:282)\n",
      "2025-06-30 01:37:11.990 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-06-30 01:37:12.003 | info     |   SiliconDriver | Opening local chip ids/pci ids: {0}/[0] and remote chip ids {} (cluster.cpp:147)\n",
      "2025-06-30 01:37:12.007 | info     |   SiliconDriver | Software version 6.0.0, Ethernet FW version 6.15.0 (Device 0) (cluster.cpp:1039)\n",
      "2025-06-30 01:37:12.039 | info     |           Metal | AI CLK for device 0 is:   1000 MHz (metal_context.cpp:122)\n",
      "2025-06-30 01:37:12.338 | info     |           Metal | Initializing device 0. Program cache is enabled (device.cpp:429)\n"
     ]
    }
   ],
   "source": [
    "import ttnn\n",
    "\n",
    "device_id = 0\n",
    "device = ttnn.open_device(device_id=device_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enable program cache\n",
    "\n",
    "Enabling the program cache will speed up the execution of operations that run repeatedly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-06-30 01:37:13.333 | info     |           Metal | Enabling program cache on MeshDevice 1 (mesh_device.cpp:537)\n"
     ]
    }
   ],
   "source": [
    "device.enable_program_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 1024\n",
    "k = 1024\n",
    "n = 1024"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize tensors a and b with random values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = ttnn.rand((m, k), dtype=ttnn.bfloat16, device=device, layout=ttnn.TILE_LAYOUT)\n",
    "b = ttnn.rand((k, n), dtype=ttnn.bfloat16, device=device, layout=ttnn.TILE_LAYOUT)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix multiply tensor a and b\n",
    "The operation will run longer the first time because the kernels need to get compiled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = a @ b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Re-running the operation shows significant speed up by utilizing program caching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "output = a @ b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the layout of matrix multiplication output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layout.TILE\n"
     ]
    }
   ],
   "source": [
    "print(output.layout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, matrix multiplication produces outputs in a tile layout. That is because it's much more efficient to use this layout for computing matrix multiplications on Tenstorrent accelerators compared to a row-major layout.\n",
    "\n",
    "And this is also why the logs show 2 tilize operations, as the inputs get automatically convered to the tile layout if they are in a row-major layout.\n",
    "\n",
    "Learn more about tile layout [here](https://github.com/tenstorrent/tt-metal/blob/main/tech_reports/tensor_layouts/tensor_layouts.md#32-tiled-layout)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the result of the matrix multiplication\n",
    "\n",
    "To inspect the results we will first convert to row-major layout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Printing ttnn tensor\n",
      "shape: Shape([1024, 1024])\n",
      "chunk of a tensor:\n",
      "ttnn.Tensor([[256.00000, 260.00000,  ..., 262.00000, 268.00000]], shape=Shape([1, 32]), dtype=DataType::BFLOAT16, layout=Layout::ROW_MAJOR)\n"
     ]
    }
   ],
   "source": [
    "output = ttnn.to_layout(output, ttnn.ROW_MAJOR_LAYOUT)\n",
    "\n",
    "print(\"Printing ttnn tensor\")\n",
    "print(f\"shape: {output.shape}\")\n",
    "print(f\"chunk of a tensor:\\n{output[:1, :32]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix multiply tensor a and b by using more performant config\n",
    "By default, matrix multiplication might not be as effecient as it could be. To speed it up further, the user can specify how many cores they want matrix multiplication to use. This can speed up the operation significantly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = ttnn.rand((m, k), dtype=ttnn.bfloat16, device=device, layout=ttnn.TILE_LAYOUT, memory_config=ttnn.L1_MEMORY_CONFIG)\n",
    "b = ttnn.rand((k, n), dtype=ttnn.bfloat16, device=device, layout=ttnn.TILE_LAYOUT, memory_config=ttnn.L1_MEMORY_CONFIG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run once to compile the kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = ttnn.matmul(a, b, memory_config=ttnn.L1_MEMORY_CONFIG, core_grid=ttnn.CoreGrid(y=8, x=8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enjoy a massive speed up on the subsequent runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = ttnn.matmul(a, b, memory_config=ttnn.L1_MEMORY_CONFIG, core_grid=ttnn.CoreGrid(y=8, x=8))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Close the device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-06-30 01:37:15.650 | info     |           Metal | Closing mesh device 1 (mesh_device.cpp:488)\n",
      "2025-06-30 01:37:15.652 | info     |           Metal | Closing mesh device 0 (mesh_device.cpp:488)\n",
      "2025-06-30 01:37:15.652 | info     |           Metal | Closing device 0 (device.cpp:469)\n",
      "2025-06-30 01:37:15.652 | info     |           Metal | Disabling and clearing program cache on device 0 (device.cpp:781)\n"
     ]
    }
   ],
   "source": [
    "ttnn.close_device(device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
